/***************************************************************************\
 * The Mattise Kernel: Userland Software									*
 * Copyright 2007 Matthew Iselin											*
 * Licensed under the GPL													*
 *																			*
 * main.c																	*
 *																			*
 * Entrypoint for the ATA driver											*
 *																			*
\***************************************************************************/

/***************************************************************************
 * Includes
 ***************************************************************************/
#include <stdio.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <process.h>
#include <mmsg.h>
#include <vfs.h>
#include "ata.h"

/***************************************************************************
 * Globals
 ***************************************************************************/

// all ATA controllers
static Tata_ctl_info ata_ctl_info[] = {
   { .base = {0x1f0, 0x3f0} },
   { .base = {0x170, 0x370} },
   { .base = {0x1e8, 0x3e0} },
   { .base = {0x168, 0x360} }
};

// debug mode
// #define DEBUGMODE

/***************************************************************************
 * Functions
 ***************************************************************************/

/// Output to a port, byte size
void outportb( short port, char dat )
{
	__asm__ __volatile__ ( "outb %1,%0" : : "dN" (port), "a" (dat) );
}

/// Output to a port, short size
void outportw( short port, short dat )
{
	__asm__ __volatile__ ( "outw %1,%0" : : "dN" (port), "a" (dat) );
}

/// Output to a port, int size
void outportl( short port, int dat )
{
	__asm__ __volatile__ ( "outl %1,%0" : : "dN" (port), "a" (dat) );
}

/// Input from a port, byte size
char inportb( short port )
{
	char dat;
	__asm__ __volatile__ ( "inb %1, %0" : "=a" (dat) : "dM" (port) );
	return dat;
}

/// Input from a port, short size
short inportw( short port )
{
	short dat;
	__asm__ __volatile__ ( "inw %1, %0" : "=a" (dat) : "dM" (port) );
	return dat;
}

/// Input from a port, short size
int inportl( short port )
{
	int dat;
	__asm__ __volatile__ ( "inl %1, %0" : "=a" (dat) : "dM" (port) );
	return dat;
}

/// Prologue to any disk operation - CHS
void chs_prolog( int addr, Drive* drv )
{
	// convert to CHS first
	unsigned char Sector = ( addr % drv->ident.sectors ) + 1;
	unsigned short Cylinder = ( addr / drv->ident.sectors ) / drv->ident.heads;
	unsigned char Head = ( addr / drv->ident.sectors ) % drv->ident.heads;
	
	// null to features
	outportb( drv->base + ATA_FEATS, 0x00 );
	
	// number of sectors
	outportb( drv->base + ATA_SCOUNT, 0x01 );
	
	// block address
	outportb( drv->base + ATA_SNUM, Sector );
	outportb( drv->base + ATA_CLOW, (unsigned char) (Cylinder & 0xFF) );
	outportb( drv->base + ATA_CHI, (unsigned char) ((Cylinder & 0xFF00)>>16) );
	
	// drive id et al
	outportb( drv->base + ATA_HDEV, ( drv->drv << 4 ) | ( Head & 0x0F ) );
}

/// Prologue to any disk operation - all except the command itself (LBA28)
void lba28_prolog( int addr, Drive* drv )
{
	// wait until we're ready
	while( (inportb( drv->base + 0x206 ) & 0x80) );
	
	// null to features
	outportb( drv->base + ATA_FEATS, 0x00 );
	
	// number of sectors
	outportb( drv->base + ATA_SCOUNT, 0x01 );
	
	// block address
	outportb( drv->base + ATA_SNUM, (uchar_t) addr );
	outportb( drv->base + ATA_CLOW, (uchar_t) (addr >> 8) );
	outportb( drv->base + ATA_CHI, (uchar_t) (addr >> 16) );

	// drive id et al
	uchar_t hdev = 0xE0 | ( drv->drv << 4 ) | ( addr >> 24 );
	outportb( drv->base + ATA_HDEV, hdev );
}

/// Prologue to any disk operation - all except the command itself (LBA48)
void lba48_prolog( unsigned long long addr, Drive* drv )
{
	// nulls to features
	outportb( drv->base + ATA_FEATS, 0x00 );
	outportb( drv->base + ATA_FEATS, 0x00 );
	
	// number of sectors
	outportb( drv->base + ATA_SCOUNT, 0x00 );
	outportb( drv->base + ATA_SCOUNT, 0x01 );
	
	// block address
	outportb( drv->base + ATA_SNUM, (uchar_t) (addr >> 24) );
	outportb( drv->base + ATA_SNUM, (uchar_t) addr );
	outportb( drv->base + ATA_CLOW, (uchar_t) (addr >> 32) );
	outportb( drv->base + ATA_CLOW, (uchar_t) (addr >> 8) );
	outportb( drv->base + ATA_CHI, (uchar_t) (addr >> 40) );
	outportb( drv->base + ATA_CHI, (uchar_t) (addr >> 16) );
	
	// drive id et al
	outportb( drv->base + ATA_HDEV, 0x40 | ( drv->drv << 4 ) );
}

/// Starts a read (CHS)
void chs_startread( int lba_addr, Drive* drv )
{
	// start the op
	chs_prolog( lba_addr, drv );

	// read command
	outportb( drv->base + ATA_COMM, 0x20 );
}

/// Starts a write (CHS)
void chs_startwrite( int lba_addr, Drive* drv )
{
	// start the op
	chs_prolog( lba_addr, drv );

	// read command
	outportb( drv->base + ATA_COMM, 0x30 );
}

/// Starts a read (LBA28)
void lba28_startread( int lba_addr, Drive* drv )
{	
	// start the op
	lba28_prolog( lba_addr, drv );

	// read command
	outportb( drv->base + ATA_COMM, 0x20 );
	
	// wait until we're ready
	while( (inportb( drv->base + ATA_COMM ) & 0x80) );
}

/// Starts a write (LBA28)
void lba28_startwrite( int lba_addr, Drive* drv )
{
	// start the op
	lba28_prolog( lba_addr, drv );

	// write command
	outportb( drv->base + ATA_COMM, 0x30 );
	
	// wait until we're ready
	while( (inportb( drv->base + ATA_COMM ) & 0x80) );
}

/// Starts a read (LBA48)
void lba48_startread( unsigned long long lba_addr, Drive* drv )
{
	// start the op
	lba48_prolog( lba_addr, drv );

	// read command
	outportb( drv->base + ATA_COMM, 0x24 );
}

/// Starts a write (LBA48)
void lba48_startwrite( unsigned long long lba_addr, Drive* drv )
{
	// start the op
	lba48_prolog( lba_addr, drv );

	// read command
	outportb( drv->base + ATA_COMM, 0x34 );
}
	
// for debugging purposes
#define ATA_LBA28_RW_TYPE	ushort_t
#define ATA_LBA28_INPORT( d ) inportw( d )
#define ATA_LBA28_OUTPORT( p, d ) outportw( p, d )

/// Reads a sector of data (LBA28)
void lba28_ReadSector( int drv, int block, char* dat, Drive* drive )
{
	// start the read
	lba28_startread( block, drive );
	
	// check for validity
	uchar_t errreg = inportb( drive->base + 0x206 ) & 0x1;
	
	// die if there was an error
	if( errreg )
	{
		return;
	}
	
	// wait until the request is done
	while( ! ( inportb( drive->base + 0x206 ) & 0x08 ) ) {};
	
	// counter
	int idx;
	
	// read in the data
	memset( dat, 0, 512 );
	for( idx = 0; idx < drive->ident.sector_bytes; idx += sizeof( ushort_t ) )
	{
		ushort_t tmpword = inportw( drive->base );
		*((ushort_t*) &dat[idx]) = tmpword;
	}
}

/// Write a sector of data (LBA28)
void lba28_WriteSector( int drv, int block, char* dat, Drive* drive )
{
	// start the write
	lba28_startwrite( block, drive );
	
	// check for validity
	uchar_t errreg = inportb( drive->base + 0x206 ) & 0x1;
	
	// die if there was an error
	if( errreg )
	{
		return;
	}
	
	// wait until the request is done
	while( ! ( inportb( drive->base + 0x206 ) & 0x08 ) ) {};
	
	// counter
	int idx;
	for( idx = 0; idx < drive->ident.sector_bytes; idx += sizeof( ushort_t ) )
	{
		outportw( drive->base, *((ushort_t*)&dat[idx]) );
	}
}

/// Read a sector of data (LBA48)
void lba48_ReadSector( int drv, unsigned long long block, char* dat, Drive* drive )
{	
	// start the read
	lba48_startread( block, drive );
	
	// check for validity
	uchar_t errreg = inportb( drive->base + ATA_ERR );
	
	// die if there was an error
	if( errreg )
	{
		return;
	}
	
	// wait for bytes
	while( ! ( inportb( drive->base + 7 ) & 0x08 ) ) {};
	
	// counter
	int idx;
	for( idx = 0; idx < 512; idx += sizeof( short ) )
	{
		short tmpword = inportw( drive->base );
		*((short*) &dat[idx]) = tmpword;
	}
}

/// Write a sector of data (LBA48)
void lba48_WriteSector( int drv, unsigned long long block, char* dat, Drive* drive )
{	
	// start the read
	lba48_startwrite( block, drive );
	
	// check for validity
	uchar_t errreg = inportb( drive->base + ATA_ERR );
	
	// die if there was an error
	if( errreg )
	{
		return;
	}
	
	// wait for bytes
	while( ! ( inportb( drive->base + 7 ) & 0x08 ) ) {};
	
	// counter
	int idx;
	for( idx = 0; idx < 512; idx += sizeof( short ) )
	{
		outportw( drive->base, *((short*)&dat[idx]) );
	}
}

/// Read a sector of data (CHS)
void chs_ReadSector( int drv, unsigned long long block, char* dat, Drive* drive )
{	
	// start the read
	chs_startread( block, drive );
	
	// check for validity
	uchar_t errreg = inportb( drive->base + ATA_ERR );
	
	// die if there was an error
	if( errreg )
	{
		return;
	}
	
	// wait for bytes
	while( ! ( inportb( drive->base + 7 ) & 0x08 ) ) {};
	
	// counter
	int idx;
	for( idx = 0; idx < 512; idx += sizeof( short ) )
	{
		short tmpword = inportw( drive->base );
		*((short*) &dat[idx]) = tmpword;
	}
}

/// Write a sector of data (CHS)
void chs_WriteSector( int drv, int block, char* dat, Drive* drive )
{	
	// start the read
	chs_startwrite( block, drive );
	
	// check for validity
	uchar_t errreg = inportb( drive->base + ATA_ERR );
	
	// die if there was an error
	if( errreg )
	{
		return;
	}
	
	// wait for bytes
	while( ! ( inportb( drive->base + 7 ) & 0x08 ) ) {};
	
	// counter
	int idx;
	for( idx = 0; idx < 512; idx += sizeof( short ) )
	{
		outportw( drive->base, *((short*)&dat[idx]) );
	}
}

/// Executes the 'Identify Device' command
IDENTIFYDEVICE IdentifyDevice( int base, int drv )
{	
	// choose the drive
	outportb( base + ATA_HDEV, 0xE0 | drv );
	
	// send & wait
	outportb( base + ATA_COMM, 0xEC );
	
	// wait until we're ready
	while( (inportb( base + 0x206 ) & 0x80) );
	
	// read in the status byte
	unsigned char ident_stat = inportb( base + ATA_COMM );
	
	// error?
	if( ident_stat & 0x1 )
	{
		// TODO: handle packet device identification, page 105 of the ATA specs
		return;
	}
	
	// return data
	IDENTIFYDEVICE ident;
	
	// wait for the data to arrive - but timeout if it doesn't
	if( ( inportb( base + ATA_COMM ) & 0x08 ) )
	{
		// counter
		int idx;
		
		// read in all the data
		unsigned short* identdata = (unsigned short*) &ident;
		for( idx = 0; idx < 255; idx++ )
		{
			identdata[idx] = inportw( base );
		}
	}

	// return it all
	return ident;
}

// ATA handler - handles I/O operations on a drive
void ata_handler()
{
	// message size and message buffer
	int msize;
	char* mbuff;
	
	// already, we will have a message waiting (INITTHRD)
	struct RecvMsgData myMsg;
	while( !mattiseHasMessage() );
	
	// allocate a buffer for the message
	msize = mattiseGetMessageSize();
	mbuff = mattiseAllocMessage();
	
	// get the message
	mattiseRecvMessage( mbuff, msize + 1, &myMsg );
	
	// open the message, get the data
	struct vfs_ReqData req;
	struct vfs_ReqData* init = (struct vfs_ReqData*) mbuff;
	
	// take in the info
	int drvid = init->reqArg1;
	int mybase = init->reqArg2;
	int myprocess = init->reqArg3;
	
	// free the memory
	mattiseFreeMessage( mbuff );

	// setup our internal information
	Drive mydrvinfo;
	mydrvinfo.base = mybase;
	mydrvinfo.drv = drvid;
	mydrvinfo.ident = IdentifyDevice( mybase, drvid );
	
	// detection
	int lba = mydrvinfo.ident.capability & 0x2;
	if( lba )
		lba = 1; // TODO: make it work with lba48 as well
	
	// internal address
	unsigned long long myInternalOffset = 0;
	
	// return the descriptor
	req.descriptor = getpid();
	mattiseSendMessage( myprocess, VFS_REPMSG, (char*) &req, sizeof( struct vfs_ReqData ) );
	
	// and now enter our message loop
	int running = 1;
	int i;
	while( running )
	{
		// do we have a message?
		if( mattiseHasMessage() )
		{
			// allocate a buffer for the message
			msize = mattiseGetMessageSize();
			mbuff = mattiseAllocMessage();
			
			// verify the buffer
			if( !mbuff )
				continue;
			
			// get the message
			mattiseRecvMessage( mbuff, msize + 1, &myMsg );
			
			// get the message
			struct vfs_ReqData* req = (struct vfs_ReqData*) mbuff;
			
			// what was it?
			switch( myMsg.msgcode )
			{
				// read: send back some data
				case VFS_READ:
				
#ifdef DEBUGMODE
					printf( "[ATA_HANDLER=%d] VFS_READ [%d]\n", getpid(), req->srcpid );
#endif
					
					// reply
					for( i = 0; i < (req->datlen); i++ )
					{
						// are we LBA48?
						if( lba == 2 )
							lba48_ReadSector( drvid, myInternalOffset + i, req->data + (i*512), &mydrvinfo );
						else if( lba == 1 )
							lba28_ReadSector( drvid, myInternalOffset + i, req->data + (i*512), &mydrvinfo );
						else
							chs_ReadSector( drvid, myInternalOffset + i, req->data + (i*512), &mydrvinfo );
					}
					req->datlen = i*512;
					mattiseSendMessage( req->srcpid, VFS_READ, (char*) req, sizeof( struct vfs_ReqData ) );
					
					break;
					
				// write: write some data
				case VFS_WRITE:
				
#ifdef DEBUGMODE
					printf( "[ATA_HANDLER=%d] VFS_WRITE [%d]\n", getpid(), req->srcpid );
#endif
					
					// reply
					for( i = 0; i < (req->datlen); i++ )
					{
						// are we LBA48?
						if( lba == 2 )
							lba48_WriteSector( drvid, myInternalOffset + i, req->data + (i*512), &mydrvinfo );
						else if( lba == 1 )
							lba28_WriteSector( drvid, myInternalOffset + i, req->data + (i*512), &mydrvinfo );
						else
							chs_WriteSector( drvid, myInternalOffset + i, req->data + (i*512), &mydrvinfo );
					}
					req->datlen = i*512;
					mattiseSendMessage( req->srcpid, VFS_WRITE, (char*) req, sizeof( struct vfs_ReqData ) );
					
					break;
				
				// seek: seek to a new location
				case VFS_SEEK:
				
#ifdef DEBUGMODE
					printf( "[ATA_HANDLER=%d] VFS_SEEK [%d]\n", getpid(), req->srcpid );
#endif
				
					// where do we seek to?
					switch( req->reqArg2 )
					{
						case SEEK_SET:
						
							// from the start
							myInternalOffset = req->reqArg1;
							
							break;
							
						case SEEK_CUR:
						
							// from now
							myInternalOffset += req->reqArg1;
					};
					
					// return new position
					req->reqArg1 = myInternalOffset;
					mattiseSendMessage( req->srcpid, VFS_SEEK, (char*) req, sizeof( struct vfs_ReqData ) );
					
					break;
				
				// getc/putc don't do anything
				case VFS_GETC:
				case VFS_PUTC:
				
#ifdef DEBUGMODE
					printf( "[ATA_HANDLER=%d] VFS_GETC/VFS_PUTC [%d]\n", getpid(), req->srcpid );
#endif
				
					req->datlen = 0;
					req->reqArg1 = -1;
					mattiseSendMessage( req->srcpid, req->msgcode, (char*) req, sizeof( struct vfs_ReqData ) );
					
					break;
				
				// stat: give some information about the device
				case VFS_STAT:
				
					{
				
#ifdef DEBUGMODE
					printf( "[ATA_HANDLER=%d] VFS_STAT [%d]\n", getpid(), req->srcpid );
#endif
					
					// the return structure
					struct stat sret;
					
					// fill in the device id
					sret.st_dev = getpid();
					sret.st_blksize = 512;
					sret.st_blocks = mydrvinfo.ident.lba_capacity;
					sret.st_mode = S_IFBLK;
					
					// send it off
					memcpy( req->data, (char*) &sret, sizeof( struct stat ) );
					mattiseSendMessage( req->srcpid, VFS_STAT, (char*) req, sizeof( struct vfs_ReqData ) );
					
					}
					break;
				
				// close: just finish
				case VFS_CLOSE:
				
#ifdef DEBUGMODE
					printf( "[ATA_HANDLER=%d] VFS_CLOSE [%d]\n", getpid(), req->srcpid );
#endif
					
					// send it back (don't need to do anything special)
					//req->msgcode = myMsg.msgcode;
					//mattiseSendMessage( req->srcpid, VFS_CLOSE, (char*) req, sizeof( struct vfs_ReqData ) );
					
					// all done!
					running = 0;
					break;
			}
			
			// free the buffer
			mattiseFreeMessage( mbuff );
		}
	}
	// loop forever
	while( 1 );
}

// a hard drive
void ata_drive()
{
	// message size and message buffer
	int msize;
	char* mbuff;
	
	// already, we will have a message waiting (INITTHRD)
	struct RecvMsgData myMsg;
	while( !mattiseHasMessage() );
	
	// allocate a buffer for the message
	msize = mattiseGetMessageSize();
	mbuff = mattiseAllocMessage();
	
	// get the message
	mattiseRecvMessage( mbuff, msize + 1, &myMsg );
	
	// open the message, get the data
	struct vfs_ReqData* req = (struct vfs_ReqData*) mbuff;

	// take in the info
	int drvid = req->reqArg1;
	int mybase = req->reqArg2;
	
	// free the memory
	mattiseFreeMessage( mbuff );
	
	// mount this drive in the devfs
	struct vfs_ReqData mountReq;	
	mountReq.srcpid = getpid();
	mountReq.datlen = sprintf( mountReq.data, "dr%d", drvid );
	mountReq.msgcode = VFS_MOUNT;
	
	// send off the message
	mattiseSendMessage(	2, /* todo: find the devfs in a different way */
						VFS_MOUNT,
						(char*) &mountReq,
						sizeof( struct vfs_ReqData )
						);
	
	// internal address
	unsigned long long myInternalOffset = 0;
	
	// and now enter our message loop
	int running = 1;
	int i;
	while( running )
	{
		// do we have a message?
		if( mattiseHasMessage() )
		{
			// allocate a buffer for the message
			msize = mattiseGetMessageSize();
			mbuff = mattiseAllocMessage();
			
			// verify the buffer
			if( !mbuff )
				continue;
			
			// get the message
			mattiseRecvMessage( mbuff, msize + 1, &myMsg );
			
			// get the message
			struct vfs_ReqData* req = (struct vfs_ReqData*) mbuff;
			struct vfs_ReqData* rep = (struct vfs_ReqData*) malloc( sizeof( struct vfs_ReqData ) );
			
			// what was it?
			switch( myMsg.msgcode )
			{
				// open: send back a descriptor
				case VFS_OPEN:
					
					{
				
#ifdef DEBUGMODE
						printf( "[ATA_DRIVE=%d] VFS_OPEN [%d]\n", getpid(), req->srcpid );
#endif
					
						// create a new thread
						int desc = createthread( (int) ata_handler, 2, "ata handler thread", 10 );
						
						// also send the init packet to the new thread
						rep->reqArg1 = drvid;
						rep->reqArg2 = mybase;
						rep->reqArg3 = req->srcpid;
						mattiseSendMessage( desc, VFS_REPMSG, (char*) rep, sizeof( struct vfs_ReqData ) );
					}
					break;
			}
			
			// free the memory
			free( rep );
			mattiseFreeMessage( mbuff );
		}
	}
	// loop forever
	while( 1 );
}

/// Drive detection
int DetectDrives()
{
	// setup for thread creation messages
	struct vfs_ReqData* req = (struct vfs_ReqData*) malloc( sizeof( struct vfs_ReqData ) );
	req->reqArg1 = 0; // drive id
	req->reqArg2 = 0; // I/O base
	req->descriptor = 0x000; // base
	
	// drive id, internal counters
	int drvi, x, y, z, intdrive = 0;
	
	// loop through each possible drive and controller
	for( drvi = 0, x = 0; x < ( sizeof( ata_ctl_info) / sizeof( Tata_ctl_info ) ); x++ )
	{
		// each controller (pcmattman: I think this is right, kmcguire would know best)
		for( z = 0; z < 2; z++ )
		{
			// master/slave - so 2 (pcmattman: again, I think this is right, kmcguire would know best)
			for(y = 0; y < 2; ++y)
			{
				// output the drive and head to the drive/head port
				outportb(ata_ctl_info[x].base[z] + ATA_HDEV, ((drvi << 4) | 0xa0 | (y << 4)));
				
				// BIOS "NEWIODELAY" macro uses this as well, so it must work
				outportb( 0xEB, 0 );
				
				// uncomment the following for debugging purposes
				// printf("[ATA] ctl:%x alt:%x drv:%x (%02x)\n", x, z, y, inportb(ata_ctl_info[x].base[z] + ATA_STATUS));
				
				// BIOS "NEWIODELAY" macro uses this as well, so it must work
				outportb( 0xEB, 0 );

				// check for the ready bit, the drive seek completion bit, no write faults, and no errors
				if((((unsigned char)inportb(ata_ctl_info[x].base[z] + ATA_STAT) & (0x40 | 0x10 | 0x20 | 0x01)) == 0x50))
				{
					// setup the thread info packet
					req->reqArg1 = drvi;
					req->reqArg2 = 0x1F0;
					
					// finally create the thread
					char* nm = (char*) malloc(32);
					sprintf( nm, "dr%d", intdrive );
					int pid = createthread( (int) ata_drive, 2, nm, 10 );
					free( nm );
					
					// send off the message to it
					mattiseSendMessage( pid, VFS_INITTHRD, (char*) req, sizeof( struct vfs_ReqData ) );
					
					// increment our internal drive id
					intdrive++;
				}
			}
		}
	}
	free( req );
}

// Entry point
int main( int argc, char* argv[] )
{
	// detect the drives
	DetectDrives();
	
	// loop forever
	while( 1 );
	
	// return success
	return 0;
}
